#!/usr/local/bin/python

import argparse
import os
import sys
import hashlib
import socket
import psycopg2
import inspect

thisscript= inspect.getfile(inspect.currentframe())

parser = argparse.ArgumentParser(description= """

DESCRIPTION:
                                 
    Insert into table bamfiles the files passed as arg.
    Given the name of one of more files and the appropriate parameters, insert_sblab_bamfiles.py 
    will insert into table bamfiles the corresponding record.

EXAMPLE:
    
    Insert into bamfiles all the files in dir bamfiles with extension .bam. Aligner is bwa and
    reference genome hg18:
    
    ls bamfiles/*.bam | insert_sblab_bamfiles.py -f - --aligner "bwa 0.5.6" --refseq hg18
    
    Will issue to postgres a query like:
    INSERT INTO bamfiles VALUES ('file1.bam', 'thishost.org', '/home/.../bamfiles', ...);
    
""", formatter_class= argparse.RawTextHelpFormatter)

parser.add_argument('-f', '--files',
                   type= str,
                   nargs= '+',
                   required= True,
                   help= '''Files to send to postgres. use "--files -" to read the list of files from stdin.
                   
                   ''')

parser.add_argument('--refseq',
                   type= str,
                   required= True,
                   help= '''String for field 'reference_seq'. E.g. mm9 or hg18
                   
                   ''')

parser.add_argument('--aligner',
                   type= str,
                   required= True,
                   help= '''String for field 'aligner'. E.g. "bwa 0.5.8c (r1536)"
                   
                   ''')

parser.add_argument('--description',
                   type= str,
                   required= False,
                   default= 'This string inserted by %s' %(thisscript),
                   help= '''String for field 'description'. Default is a message with the name of this script
                   
                   ''')

parser.add_argument('--touse',
                   type= bool,
                   required= False,
                   default= True,
                   help= '''String for field 'to_use'. Is this bamfile representative for the corresponding library? NB: default is True"
                   
                   ''')

parser.add_argument('--connection', '-c',
                   type= str,
                   required= False,
                   help= '''A string to be passed to psycopg2.connect() to connect to the database. E.g. "dbname=testdb user=me password= mypwd host= localhost"
If None, connection parameters will be read from ~/.psycopgpass

''')

parser.add_argument('--update', '-u',
                   action= 'store_true',
                   required= False,
                   help='''Not implemented yet. If a record with name equal to input is found, update it. (Default is to throw an error due to primary key violation).

''')

args = parser.parse_args()

# --------------------------[ Functions ]---------------------------------------

def sumfile(fobj):
    '''Returns an md5 hash for an object with read() method.'''
    m = hashlib.md5()
    while True:
        d = fobj.read(8096)
        if not d:
            break
        m.update(d)
    return m.hexdigest()


def md5sum(fname):
    '''Returns an md5 hash for file fname, or stdin if fname is "-".'''
    if fname == '-':
        ret = sumfile(sys.stdin)
    else:
        try:
            f = open(fname, 'rb')
        except:
            return 'Failed to open file %s' %fname
        ret = sumfile(f)
        f.close()
    return ret
# ------------------------------------------------------------------------------

if args.connection is None:
    try:
        conn= open(os.path.join(os.getenv("HOME"), '.psycopgpass'))
    except:
        sys.exit('I cannot read file %s and -c/--connection parameter is None.\nI cannot connect to any database.' %(os.path.join(os.getenv("HOME"), '.psycopgpass')))
    conn= conn.readlines()
    conn_args= [x.strip() for x in conn if x.strip() != '' and x.strip().startswith('#') is False][0]
    print('\nNOTICE: Connecting to database using parameters in ~/.psycopgpass\n')
else:
    conn_args= args.connection

conn = psycopg2.connect(conn_args)
cur= conn.cursor()

if args.files == ['-']:
    args.files= sys.stdin.readlines()
    args.files= [x.strip() for x in args.files]
args.files= sorted(list(set(args.files)))

## Check all the files exist
for f in args.files:
    if not os.path.isfile(f):
       sys.exit('\n%s: File %s not found\n' %(thisscript, f)) 

for f in args.files:
    ## This list will be passed to target table:
    record= [os.path.split(f)[1],
             socket.gethostname(),
             os.path.split(os.path.abspath(f))[0],
             md5sum(f),
             args.refseq,
             args.aligner,
             args.description,
             args.touse
            ]
    insert_str= cur.mogrify("""INSERT INTO bamfiles VALUES (%s, %s, %s, %s, %s, %s, %s, %s)""", record )
    print(insert_str)
    cur.execute(insert_str)
conn.commit()
cur.close()
conn.close()